from langchain_community.vectorstores import Qdrant
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain_community.llms import Tongyi
from langchain_community.embeddings import QianfanEmbeddingsEndpoint
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader, PyPDFLoader
import os
import json
import warnings
warnings.filterwarnings('ignore')

# 读取API Key
with open(r'./ERNIE_config.json') as f:
    key = f.read()
    api_key = json.loads(key)

# 设置通义千问API Key
DASHSCOPE_API_KEY = api_key['DASHSCOPE_API_KEY']  # 阿里通义千问模型DASHSCOPE_API_KEY
os.environ["DASHSCOPE_API_KEY"] = DASHSCOPE_API_KEY  # 设置环境变量

# 加载PDF文档
loader = PyPDFLoader('01_"未来校园"智能应用专项赛.pdf')  # 读取数据
documents = loader.load()  # 将数据转为指定格式

# 文本分割
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)  # 指定切分时每份文件的token数
chunked_documents = text_splitter.split_documents(documents)  # 将文本切分成指定长度大小

# 使用通义千问的嵌入模型
# 注意：由于通义千问没有直接的Embeddings实现，这里使用QianfanEmbeddingsEndpoint作为替代
# 实际使用时可能需要根据最新的langchain支持情况调整
embedding = QianfanEmbeddingsEndpoint(
    model="bce-embedding-base",  # 使用百度千帆的嵌入模型作为替代
    api_key=DASHSCOPE_API_KEY
)

# 利用Embedding将文本向量化，并存储在向量数据库中
vectorstore = Qdrant.from_documents(
    documents=chunked_documents,  # 已分块的文档
    embedding=embedding,  # 使用嵌入模型
    location=":memory:",  # 存储在内存中
    collection_name="my_documents"
)  # 指定collection_name

retriever = vectorstore.as_retriever()  # 索引器

# 配置上下文（context）及用户输入（question）
setup_and_retrieval = RunnableParallel(
    {"context": retriever, "question": RunnablePassthrough()}
)

template = """Answer the question based only on the following context:
    {context}

Question: {question}
"""

prompt = ChatPromptTemplate.from_template(template)

# 配置通义千问模型
model_tongyi = Tongyi(
    model="qwen-plus",  # 使用通义千问Plus模型
    temperature=0.7,  # 控制生成文本的随机性
    api_key=DASHSCOPE_API_KEY
)

# 输出解析器
output_parser = StrOutputParser()

# 构建完整的链（Chain）
chain_tongyi = setup_and_retrieval | prompt | model_tongyi | output_parser

# 效果测试
def ask_question(question):
    """向RAG系统提问并获取回答"""
    answer = chain_tongyi.invoke(question)
    return answer

# 示例使用
if __name__ == "__main__":
    # 测试问题
    test_questions = [
        "未来校园智能应用专项赛的主要内容是什么？",
        "这个比赛的参赛对象有哪些？"
    ]
    
    for question in test_questions:
        print(f"问题: {question}")
        answer = ask_question(question)
        print(f"回答: {answer}")
        print("-" * 50)